{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Modeling and Simulation in Python\n",
    "\n",
    "Chapter 11\n",
    "\n",
    "Copyright 2017 Allen Downey\n",
    "\n",
    "License: [Creative Commons Attribution 4.0 International](https://creativecommons.org/licenses/by/4.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure Jupyter so figures appear in the notebook\n",
    "%matplotlib inline\n",
    "\n",
    "# Configure Jupyter to display the assigned value after an assignment\n",
    "%config InteractiveShell.ast_node_interactivity='last_expr_or_assign'\n",
    "\n",
    "# import functions from the modsim.py module\n",
    "from modsim import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SIR implementation\n",
    "\n",
    "We'll use a `State` object to represent the number (or fraction) of people in each compartment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "init = State(S=89, I=1, R=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To convert from number of people to fractions, we divide through by the total."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "init /= sum(init)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`make_system` creates a `System` object with the given parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_system(beta, gamma):\n",
    "    \"\"\"Make a system object for the SIR model.\n",
    "    \n",
    "    beta: contact rate in days\n",
    "    gamma: recovery rate in days\n",
    "    \n",
    "    returns: System object\n",
    "    \"\"\"\n",
    "    init = State(S=89, I=1, R=0)\n",
    "    init /= sum(init)\n",
    "\n",
    "    t0 = 0\n",
    "    t_end = 7 * 14\n",
    "\n",
    "    return System(init=init, t0=t0, t_end=t_end,\n",
    "                  beta=beta, gamma=gamma)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's an example with hypothetical values for `beta` and `gamma`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "tc = 3      # time between contacts in days \n",
    "tr = 4      # recovery time in days\n",
    "\n",
    "beta = 1 / tc      # contact rate in per day\n",
    "gamma = 1 / tr     # recovery rate in per day\n",
    "\n",
    "system = make_system(beta, gamma)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The update function takes the state during the current time step and returns the state during the next time step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_func(state, t, system):\n",
    "    \"\"\"Update the SIR model.\n",
    "    \n",
    "    state: State with variables S, I, R\n",
    "    t: time step\n",
    "    system: System with beta and gamma\n",
    "    \n",
    "    returns: State object\n",
    "    \"\"\"\n",
    "    s, i, r = state\n",
    "\n",
    "    infected = system.beta * i * s    \n",
    "    recovered = system.gamma * i\n",
    "    \n",
    "    s -= infected\n",
    "    i += infected - recovered\n",
    "    r += recovered\n",
    "    \n",
    "    return State(S=s, I=i, R=r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run a single time step, we call it like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "state = update_func(init, 0, system)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can run a simulation by calling the update function for each time step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_simulation(system, update_func):\n",
    "    \"\"\"Runs a simulation of the system.\n",
    "    \n",
    "    system: System object\n",
    "    update_func: function that updates state\n",
    "    \n",
    "    returns: State object for final state\n",
    "    \"\"\"\n",
    "    state = system.init\n",
    "    \n",
    "    for t in linrange(system.t0, system.t_end):\n",
    "        state = update_func(state, t, system)\n",
    "        \n",
    "    return state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The result is the state of the system at `t_end`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "run_simulation(system, update_func)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise**  Suppose the time between contacts is 4 days and the recovery time is 5 days.  After 14 weeks, how many students, total, have been infected?\n",
    "\n",
    "Hint: what is the change in `S` between the beginning and the end of the simulation?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using TimeSeries objects"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we want to store the state of the system at each time step, we can use one `TimeSeries` object for each state variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_simulation(system, update_func):\n",
    "    \"\"\"Runs a simulation of the system.\n",
    "    \n",
    "    Add three Series objects to the System: S, I, R\n",
    "    \n",
    "    system: System object\n",
    "    update_func: function that updates state\n",
    "    \"\"\"\n",
    "    S = TimeSeries()\n",
    "    I = TimeSeries()\n",
    "    R = TimeSeries()\n",
    "\n",
    "    state = system.init\n",
    "    t0 = system.t0\n",
    "    S[t0], I[t0], R[t0] = state\n",
    "    \n",
    "    for t in linrange(system.t0, system.t_end):\n",
    "        state = update_func(state, t, system)\n",
    "        S[t+1], I[t+1], R[t+1] = state\n",
    "    \n",
    "    return S, I, R"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's how we call it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "tc = 3      # time between contacts in days \n",
    "tr = 4      # recovery time in days\n",
    "\n",
    "beta = 1 / tc      # contact rate in per day\n",
    "gamma = 1 / tr     # recovery rate in per day\n",
    "\n",
    "system = make_system(beta, gamma)\n",
    "S, I, R = run_simulation(system, update_func)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And then we can plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_results(S, I, R):\n",
    "    \"\"\"Plot the results of a SIR model.\n",
    "    \n",
    "    S: TimeSeries\n",
    "    I: TimeSeries\n",
    "    R: TimeSeries\n",
    "    \"\"\"\n",
    "    plot(S, '--', label='Susceptible')\n",
    "    plot(I, '-', label='Infected')\n",
    "    plot(R, ':', label='Recovered')\n",
    "    decorate(xlabel='Time (days)',\n",
    "             ylabel='Fraction of population')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's what they look like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_results(S, I, R)\n",
    "savefig('figs/chap11-fig01.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using a DataFrame"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instead of making three `TimeSeries` objects, we can use one `DataFrame`.\n",
    "\n",
    "We have to use `row` to selects rows, rather than columns.  But then Pandas does the right thing, matching up the state variables with the columns of the `DataFrame`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_simulation(system, update_func):\n",
    "    \"\"\"Runs a simulation of the system.\n",
    "        \n",
    "    system: System object\n",
    "    update_func: function that updates state\n",
    "    \n",
    "    returns: TimeFrame\n",
    "    \"\"\"\n",
    "    frame = TimeFrame(columns=system.init.index)\n",
    "    frame.row[system.t0] = system.init\n",
    "    \n",
    "    for t in linrange(system.t0, system.t_end):\n",
    "        frame.row[t+1] = update_func(frame.row[t], t, system)\n",
    "    \n",
    "    return frame"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's how we run it, and what the result looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "tc = 3      # time between contacts in days \n",
    "tr = 4      # recovery time in days\n",
    "\n",
    "beta = 1 / tc      # contact rate in per day\n",
    "gamma = 1 / tr     # recovery rate in per day\n",
    "\n",
    "system = make_system(beta, gamma)\n",
    "results = run_simulation(system, update_func)\n",
    "results.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can extract the results and plot them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_results(results.S, results.I, results.R)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises\n",
    "\n",
    "**Exercise**  Suppose the time between contacts is 4 days and the recovery time is 5 days.  Simulate this scenario for 14 weeks and plot the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution goes here"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
